%This script runs the simulations used for Figure B3
%Corresponding author: Rodrigo Adao
%Date: 09/11/2024
%Input: model_data_fgkk.mat
%Output: Figures B3

%% Preliminaries

clear all;
close all;

%Define vector of large varieties in our sample
trade_target = 0.9; %baseline has 0.9

%Simulation parameters
Nsimul = 1500; %Number of simulations (baseline 1500)

%% Import Data
%set local paths
function_path = "..\Functions\";
addpath(function_path,'-frozen');
set_local_paths;

%Import data and compute researcher predictions
simulation_preliminary_steps;

%Numerical parameters
tol_conv = 1e-6;    %tolerance level to check for price convergence
adj_inner_g0 = .4;

%% Create IV matrices 

 %Naive IV matrices
    var_adj=sum(ww_dW.^2)/Nobs;
    share_NC_nGE  = shareIV_tau;
    share_wNC_nGE  = ww_dWmat_naive*share_NC_nGE; 
    share_wMC_nGE  = ( share_wNC_nGE - Cn*(MCn_term1*share_wNC_nGE) )/var_adj;
    clearvars share_wNC_nGE share_NC_nGE shareIV_tau  

%Auxiliary matrices for control and welfare adjustments
%load auxiliary adjustment vectors
load(save_share_path + "simulation_shares_t" + trade_target + ".mat", 'adjGE_wMC')

    adjGE_wMC_mat= spdiags(adjGE_wMC, 0 , Nobs, Nobs );
    share_wMC_dWa = adjGE_wMC_mat*shareMOD_dW;
    share_wMC_dW = share_wMC_dWa - C*(MC_term1*share_wMC_dWa);
    clearvars share_wMC_dWa adjGE_wMC_mat ww_dWmat ww_dWmat_naive

%% Compute equilibrium given different draws of parameters

%Parameters for simulation draws
sgrid = [.02; .04; .06; .08; .10];
Ns = length(sgrid);

sd_shifters_grid   = sgrid;
mean_shifters_grid = sgrid;

%parameters for draws of shocks to other parameters from normal distribution
mean_deta_grid     = 0*ones(1*Ns,1); 
sd_da_grid         = sgrid;  
sd_dzstar_grid     = sgrid;  
sd_dastar_grid     = sgrid; 

%Parameters controlling true DGP
sigma_grid      = sigma     *ones(1*Ns,1);       
omega_star_grid = omega_star*ones(1*Ns,1);  
eta_grid        = eta       *ones(1*Ns,1);        
kappa_grid      = kappa     *ones(1*Ns,1);       
sigma_star_grid = sigma_star*ones(1*Ns,1);
rho_exp_grid    = 1         *ones(1*Ns,1);
rho_imp_grid    = 1         *ones(1*Ns,1);
gammaM_grid     = 1         *ones(1*Ns,1);
gammaX_grid     = 1         *ones(1*Ns,1);
gammaQ_grid     = 1         *ones(1*Ns,1);

gamma_mean_px_grid = 0*ones(1*Ns,1);
gamma_mean_pm_grid = 0*ones(1*Ns,1);
gamma_mean_rm_grid = 0*ones(1*Ns,1);

gamma_sd_px_grid = zeros(size(gamma_mean_px_grid));
gamma_sd_pm_grid = zeros(size(gamma_mean_pm_grid));
gamma_sd_rm_grid = zeros(size(gamma_mean_rm_grid));

Npar = Ns;

%% Compute equilibrium given different draws of parameters

Npar
ests = zeros(21, Nsimul, Npar);

for j=1:Npar
    tic
    display('------------- running new j -----------------')
    j
    
    %shock process
    sd_shifters  = sd_shifters_grid(j);
    mean_shifters = mean_shifters_grid(j);
    adj = Nobs*(mean_shifters/sd_shifters);
    %adj = Nobs;

    mean_deta = mean_deta_grid(j);
    sd_da     = sd_da_grid(j);
    sd_dzstar = sd_dzstar_grid(j);
    sd_dastar = sd_dastar_grid(j);

    %Auxiliary matrices for computing IV       
    share_wMC_dW_j = adj*share_wMC_dW;
    
    %Generate true DGP
    omega_starDGP = omega_star_grid(j);    
    sigmaDGP = sigma_grid(j);           
    etaDGP = eta_grid(j);               
    kappaDGP = kappa_grid(j);           
    sigma_starDGP = sigma_star_grid(j); 
    rhoMDGP = rho_imp_grid(j) ;
    rhoXDGP = rho_exp_grid(j) ;
    gammaMDGP = gammaM_grid(j) ;
    gammaXDGP = gammaX_grid(j) ;
    gammaQDGP = gammaQ_grid(j) ;    

    gamma_n = (indM==1).*( normrnd(gamma_mean_pm_grid(j), gamma_sd_pm_grid(j), Nobs, 1) )... 
            + (indT==1).*( normrnd(gamma_mean_rm_grid(j), gamma_sd_rm_grid(j), Nobs, 1) )...
            + (indX==1).*( normrnd(gamma_mean_px_grid(j), gamma_sd_px_grid(j), Nobs, 1) );
    gammaDGP = spdiags(1+ gamma_n, 0, Nobs,Nobs);

     [delta_ig_DGP0, a_star_ig_DGP0, z_star_ig_DGP0, a_ig_DGP0, aM_g_DGP0, AM_s_DGP0, AD_s_DGP0, betaT_DGP0, beta_s_DGP0, alphaL_s_DGP0, alphaI_s_DGP0, alpha_ks_DGP0, barL_rs_DGP0, Z_rs_DGP0, D_DGP0, E_s_DGP0, F_DGP0] = ... 
        invert_param_full(x_ig_0, m_ig_0, tau_star_ig_0, tau_ig_0, trade_con_share, tot_labor_comp, tot_intermediate, sales_s, IO_sales, Lsr, ...
        Dsg, DX, DM, omega_starDGP, sigmaDGP, etaDGP, kappaDGP, sigma_starDGP, gammaMDGP, gammaXDGP,  tol_parm);
 
    Lr = sum(Lsr,2) + LrNT';
    br_DGP0 = Lr/sum(Lr);

    [p_s_IV, PM_s_IV, ...
         rhoDGP_qm_ig_tau_ig, rhoDGP_qm_ig_taustar_ig, rhoDGP_qm_ig_a_ig, rhoDGP_qm_ig_zstar_ig, rhoDGP_qm_ig_astar_ig, rhoDGP_qm_ig_delta_ig, ...
         rhoDGP_pm_ig_tau_ig, rhoDGP_pm_ig_taustar_ig, rhoDGP_pm_ig_a_ig, rhoDGP_pm_ig_zstar_ig, rhoDGP_pm_ig_astar_ig, rhoDGP_pm_ig_delta_ig, ...
         rhoDGP_px_ig_tau_ig, rhoDGP_px_ig_taustar_ig, rhoDGP_px_ig_a_ig, rhoDGP_px_ig_zstar_ig, rhoDGP_px_ig_astar_ig, rhoDGP_px_ig_delta_ig, ...
         rhoDGP_pstar_ig_tau_ig, rhoDGP_pstar_ig_taustar_ig, rhoDGP_pstar_ig_a_ig, rhoDGP_pstar_ig_zstar_ig, rhoDGP_pstar_ig_astar_ig, rhoDGP_pstar_ig_delta_ig, ...
         rhoDGP_rm_ig_tau_ig, rhoDGP_rm_ig_taustar_ig, rhoDGP_rm_ig_a_ig, rhoDGP_rm_ig_zstar_ig, rhoDGP_rm_ig_astar_ig, rhoDGP_rm_ig_delta_ig] ...
         = compute_equilibrium_foa_full( delta_ig_DGP0, a_star_ig_DGP0, z_star_ig_DGP0, a_ig_DGP0, aM_g_DGP0, AM_s_DGP0, AD_s_DGP0, betaT_DGP0, beta_s_DGP0, alphaL_s_DGP0, alphaI_s_DGP0, alpha_ks_DGP0, barL_rs_DGP0, Z_rs_DGP0, D_DGP0, ... 
             tau_ig_0, tau_star_ig_0, m_ig_0, E_s_DGP0, p_s_0, Dsg, DX, DM, omega_starDGP, sigmaDGP, etaDGP, kappaDGP, sigma_starDGP, nu, epsilon, gammaQDGP, gammaXDGP, gammaMDGP, rhoXDGP, rhoMDGP, tol_conv, adj_inner_g0, adj_outter_g0, ... 
             ind_Mshifts_vec, ind_Xshifts_vec, ind_Msample_ig, ind_Xsample_ig ); 

    shareDGP_pm      = [rhoDGP_pm_ig_tau_ig    , rhoDGP_pm_ig_taustar_ig   ];
    shareDGP_rm      = [rhoDGP_rm_ig_tau_ig    , rhoDGP_rm_ig_taustar_ig   ];
    shareDGP_px      = [rhoDGP_px_ig_tau_ig    , rhoDGP_px_ig_taustar_ig   ];
    shareDGP_dW = gammaDGP*[shareDGP_pm; shareDGP_rm; shareDGP_px];
    
    shareDGP_dW_da     = [rhoDGP_pm_ig_a_ig    ; rhoDGP_rm_ig_a_ig    ; rhoDGP_px_ig_a_ig    ];
    shareDGP_dW_dzstar = [rhoDGP_pm_ig_zstar_ig; rhoDGP_rm_ig_zstar_ig; rhoDGP_px_ig_zstar_ig];
    shareDGP_dW_dastar = [rhoDGP_pm_ig_astar_ig; rhoDGP_rm_ig_astar_ig; rhoDGP_px_ig_astar_ig];
   
    [p_s_DGP0, PM_s_DGP0, pM_ig_DGP0, RevM_ig_DGP0, pX_ig_DGP0, VM_ig_DGP0, F_r_DGP0, RW_r_DGP0] ...   
         = compute_equilibrium_nonlinear( delta_ig_DGP0, a_star_ig_DGP0, z_star_ig_DGP0, a_ig_DGP0, aM_g_DGP0, AM_s_DGP0, AD_s_DGP0, betaT_DGP0, beta_s_DGP0, alphaL_s_DGP0, alphaI_s_DGP0, alpha_ks_DGP0, barL_rs_DGP0, Z_rs_DGP0, D_DGP0, br_DGP0, ... 
             tau_ig_0, tau_star_ig_0, m_ig_0, E_s_DGP0, p_s_0, Dsg, DX, DM, omega_starDGP, sigmaDGP, etaDGP, kappaDGP, sigma_starDGP, nu, epsilon, gammaQDGP, gammaXDGP, gammaMDGP, rhoXDGP, rhoMDGP, tol_conv, adj_inner_g0, adj_outter_g0, ... 
             ind_Mshifts_vec, ind_Xshifts_vec, ind_Msample_ig, ind_Xsample_ig );    

    clear    rhoDGP* shareDGP_p* shareDGP_r* shareDGP_q* 

    toc
    tic
    %Simulation across draws of shocks 
    parfor n=1:Nsimul
    %for n=1:Nsimul
    %    n

    % simulated data
        %simulated shifts
        dtau_sim = normrnd(mean_shifters, sd_shifters, NMs, 1);
        dtaustar_sim = normrnd(mean_shifters, sd_shifters, NXs, 1);  
        shifters = [dtau_sim; dtaustar_sim];
        shiftersIV = (shifters - mean(shifters))/std(shifters);
        
        %simulated shocks
        da_sim     = normrnd(mean_deta, sd_da    , NMs, 1);
        dzstar_sim = normrnd(mean_deta, sd_dzstar, NMs, 1);
        dastar_sim = normrnd(mean_deta, sd_dastar, NXs, 1);
        detaDGP = shareDGP_dW_da*da_sim + shareDGP_dW_dzstar*dzstar_sim + shareDGP_dW_dastar*dastar_sim ;
        
        %New level of shcoked variables for nonlinear solution
        tau_sim = reshape(tau_ig_0', N*G, 1);
        tau_sim(ind_Mshifts_vec==1) = -1 + ( 1+tau_sim(ind_Mshifts_vec==1) ).*exp( dtau_sim./(1+tau_sim(ind_Mshifts_vec==1)) );
        tau_sim = reshape(tau_sim, G, N)';
        tau_sim = max(tau_sim, -.9);
        
        taustar_sim = reshape(tau_star_ig_0', N*G, 1);
        %taustar_sim(ind_Xshifts_vec==1) = taustar_sim(ind_Xshifts_vec==1) + dtaustar_sim;
        taustar_sim(ind_Xshifts_vec==1) = -1 + ( 1+taustar_sim(ind_Xshifts_vec==1) ).*exp( dtaustar_sim./(1+taustar_sim(ind_Xshifts_vec==1)) );
        taustar_sim = reshape(taustar_sim, G, N)';
        taustar_sim = max(taustar_sim, -.9);
        
        a_sim = reshape(a_ig_DGP0', N*G, 1);
        a_sim(ind_Mshifts_vec==1) = a_sim(ind_Mshifts_vec==1).*exp(da_sim);
        a_sim = reshape(a_sim, G, N)';
        
        a_star_sim = reshape(a_star_ig_DGP0', N*G, 1);
        a_star_sim(ind_Xshifts_vec==1) = a_star_sim(ind_Xshifts_vec==1).*exp(dastar_sim);
        a_star_sim = reshape(a_star_sim, G, N)';
        
        z_star_sim = reshape(z_star_ig_DGP0', N*G, 1);
        z_star_sim(ind_Mshifts_vec==1) = z_star_sim(ind_Mshifts_vec==1).*exp(dzstar_sim);
        z_star_sim = reshape(z_star_sim, G, N)';  
        
        %Nonlinear solution for policy impact
        [p_s_sim0, PM_s_sim0, pM_ig_sim0, RevM_ig_sim0, pX_ig_sim0, VM_ig_sim0, F_r_sim0, RW_r_sim0] ...   
         = compute_equilibrium_nonlinear( delta_ig_DGP0, a_star_sim, z_star_sim, a_sim, aM_g_DGP0, AM_s_DGP0, AD_s_DGP0, betaT_DGP0, beta_s_DGP0, alphaL_s_DGP0, alphaI_s_DGP0, alpha_ks_DGP0, barL_rs_DGP0, Z_rs_DGP0, D_DGP0, br_DGP0, ... 
             tau_ig_0, tau_star_ig_0, m_ig_0, E_s_DGP0, p_s_0, Dsg, DX, DM, omega_starDGP, sigmaDGP, etaDGP, kappaDGP, sigma_starDGP, nu, epsilon, gammaQDGP, gammaXDGP, gammaMDGP, rhoXDGP, rhoMDGP, tol_conv, adj_inner_g0, adj_outter_g0, ... 
             ind_Mshifts_vec, ind_Xshifts_vec, ind_Msample_ig, ind_Xsample_ig ); 
     
        [p_s_sim1, PM_s_sim1, pM_ig_sim1, RevM_ig_sim1, pX_ig_sim1, VM_ig_sim1, F_r_sim1, RW_r_sim1] ...    
         = compute_equilibrium_nonlinear( delta_ig_DGP0, a_star_sim, z_star_sim, a_sim, aM_g_DGP0, AM_s_DGP0, AD_s_DGP0, betaT_DGP0, beta_s_DGP0, alphaL_s_DGP0, alphaI_s_DGP0, alpha_ks_DGP0, barL_rs_DGP0, Z_rs_DGP0, D_DGP0, br_DGP0, ... 
             tau_sim, taustar_sim, m_ig_0, E_s_DGP0, p_s_sim0, Dsg, DX, DM, omega_starDGP, sigmaDGP, etaDGP, kappaDGP, sigma_starDGP, nu, epsilon, gammaQDGP, gammaXDGP, gammaMDGP, rhoXDGP, rhoMDGP, tol_conv, adj_inner_g0, adj_outter_g0, ... 
             ind_Mshifts_vec, ind_Xshifts_vec, ind_Msample_ig, ind_Xsample_ig ); 
     
        deta_pM_ig   = log(pM_ig_sim0  ./pM_ig_DGP0  );
        deta_RevM_ig = (RevM_ig_sim0 - RevM_ig_DGP0)./VM_ig_DGP0;
        deta_pX_ig   = log(pX_ig_sim0  ./pX_ig_DGP0  );
        
        dx_pM_ig     = log(pM_ig_sim1  ./pM_ig_sim0  );
        dx_RevM_ig   = (RevM_ig_sim1- RevM_ig_sim0)./VM_ig_DGP0;
        dx_pX_ig     = log(pX_ig_sim1  ./pX_ig_sim0  );
        
        dx_nl   = [dx_pM_ig  ; dx_RevM_ig   ; dx_pX_ig   ];
        deta_nl = [deta_pM_ig; deta_RevM_ig; deta_pX_ig];
        
        %Simulated predictions and outcomes with true parameters
        dx     = shareMOD_dW*shifters;
        dxstar = shareDGP_dW*shifters ;
        dxstar_nl = gammaDGP*dx_nl;
        
        dy_nl = dxstar_nl + deta_nl;
        dy = dxstar + detaDGP;
        
        %data linear, model linear
        dyt_NC = dy - dx;
        dyn_NC = dyt_NC(sample_naive);
        
        %data nonlinear, model linear
        dynlt_NC = dy_nl - dx;
        dynln_NC = dynlt_NC(sample_naive);
        
        %data nonlinear, model nonlinear
        dynllt_NC = dy_nl - dx_nl;
        dynlln_NC = dynllt_NC(sample_naive);

        % Welfare statistics
        dW = ww_dW'*dx;    
        dWstar = ww_dW'*dxstar;
        Delta_W = dWstar - dW;
        
        dWstar_nl = ww_dW'*dxstar_nl;
        Delta_W_nl = dWstar_nl - dWstar;
        
        %Weighted avg of real income change times 100 (since ww_dW also multiplied by 100)
        dWstar_et = 100*(F_r_DGP0/sum(F_r_DGP0))'*log(RW_r_sim1./RW_r_sim0);
        Delta_W_et = dWstar_et - dWstar;
        
      %checking approximation
      st_dx = std([dxstar,dxstar_nl]);
      cr_dx = corr([dxstar,dxstar_nl]);
      cr_dx = cr_dx(1,2); 
      
    % Tests with the true parameters
        
       %correlation-based tests
        correl_pred = corr([dy, dx ]);
        correl_pred = correl_pred(2,1);
 
        correl_pred_nl = corr([dy_nl, dx ]);
        correl_pred_nl = correl_pred_nl(2,1);

        correl_pred_nll = corr([dy_nl, dx_nl ]);
        correl_pred_nll = correl_pred_nll(2,1);
          
       %data generated with nonlinear model, predictions with linear model  
        zn_wMC = share_wMC_nGE*shiftersIV;          
        zw_wMC = share_wMC_dW_j*shiftersIV;  

        %naive -- no GE      
        [bt_zn_wMCnl, se_zn_wMCnl, rj_zn_wMCnl, r0_zn_wMCnl] = implement_test(dynln_NC, zn_wMC, share_wMC_nGE, shiftersIV, critical, 0);
        
        %welfare adjusted
        [bt_zw_wMCnl, se_zw_wMCnl, rj_zw_wMCnl, r0_zw_wMCnl] = implement_test(dynlt_NC, zw_wMC, share_wMC_dW_j, shiftersIV, critical, 0);     

        %Output
       est_n = [dW, dWstar, dWstar_et, dWstar_nl, Delta_W, Delta_W_et, Delta_W_nl, correl_pred, correl_pred_nl, correl_pred_nll, st_dx, cr_dx, ...
               bt_zn_wMCnl, bt_zw_wMCnl, ...
               se_zn_wMCnl, se_zw_wMCnl, ...
               rj_zn_wMCnl, rj_zw_wMCnl, ...
               r0_zn_wMCnl, r0_zw_wMCnl];
    
    % save outcomes for step n
        ests(:,n,j) = est_n';
        
    end
    toc
    save(save_simulation_path + "output_Fig_B3.mat", 'ests', 'tradestat_sample', 'sgrid' , 'mean_shifters_grid', 'sd_shifters_grid', 'sd_da_grid', 'gamma_mean_px_grid','gamma_mean_pm_grid', 'gamma_mean_rm_grid', 'Npar', 'min_imp', 'min_exp', 'Ns', 'trade_target');
end

%% Report results

load(save_simulation_path + "output_Fig_B3.mat")

ests_stat=[];
for j = 1:Npar
    ests_j = ests(:,:,j);
    
     estimates_j = ests_j(1:15,:);
     SE   = [zeros(13,1); mean(ests_j(16:17,:),2)];
     rej  = [zeros(13,1); mean(ests_j(18:19,:),2)];
     rej0 = [zeros(13,1); mean(ests_j(20:21,:),2)];

    ests_stat(:,:,j)     = [mean(estimates_j,2), std(estimates_j,1,2), SE, rej0, rej];   
end

%% Figure B3

title_size = 14;
label_size = 20;
labely_size = 16;
legend_size = 19;
marker_size = 10;
axes_size = 18;

j0 = 1;
jf = Ns;

dW_lv = 1;
dW_nlv = 3;
dW_nllv = 4;

dWl_series =  squeeze(ests_stat(dW_lv,1,j0:jf));
dWnl_series =  squeeze(ests_stat(dW_nlv,1,j0:jf));
dWnll_series =  squeeze(ests_stat(dW_nllv,1,j0:jf));

naive_nl_version = 14;
ACDIV_nl_version = 15;

%Average estimated coefficientf
figure('DefaultAxesFontSize',axes_size, 'Position', [10 10 600 600]);
plot_lines_avgA = tiledlayout(1,1);

% Difference plotted on the x-axis
x_axis = dWnll_series - dWl_series;

beta_naiveIV_nl_series =  squeeze(ests_stat(naive_nl_version,1,j0:jf));
beta_ACDIV_nl_series   =  squeeze(ests_stat(ACDIV_nl_version,1,j0:jf));

ymin = min(min(beta_naiveIV_nl_series), min(beta_ACDIV_nl_series));
ymax = 0;

nexttile
plot(x_axis, beta_naiveIV_nl_series, 'Marker','d','MarkerFaceColor',[0.5 0.5 0.5], 'MarkerSize', marker_size, 'Color', [0.5 0.5 0.5])
hold on 
plot(x_axis, beta_ACDIV_nl_series  , 'Marker','o','MarkerFaceColor','black', 'MarkerSize', marker_size, 'Color', 'black')
hold on 
yline(0, '-black', 'HandleVisibility','off')
xline(0, '-black', 'HandleVisibility','off')
ylim([ymin ymax])
legend('naive IV', 'preferred IV', 'Location', 'south', 'FontSize', legend_size)
xlabel('E_t[W (\Delta x* ) - W (\Delta x )]', 'FontSize', label_size)

%Rejection rate
figure('DefaultAxesFontSize',axes_size, 'Position', [10 10 600 600]);
plot_lines_rejA = tiledlayout(1,1);

rej0_naiveIV_nl_series  =  squeeze(ests_stat(naive_nl_version,4,j0:jf));
rej0_ACDIV_nl_series    =  squeeze(ests_stat(ACDIV_nl_version,4,j0:jf));

ymin = 0 ;
ymax = max(max(rej0_naiveIV_nl_series), max(rej0_ACDIV_nl_series));

nexttile
plot(x_axis, rej0_naiveIV_nl_series, 'Marker','d','MarkerFaceColor',[0.5 0.5 0.5], 'MarkerSize', marker_size, 'Color', [0.5 0.5 0.5])
hold on 
plot(x_axis, rej0_ACDIV_nl_series  , 'Marker','o','MarkerFaceColor','black', 'MarkerSize', marker_size, 'Color', 'black')
hold on
yline(0, '-black', 'HandleVisibility','off')
xline(0, '-black', 'HandleVisibility','off')
ylim([ymin ymax])
legend('naive IV', 'preferred IV', 'Location', 'south', 'FontSize', legend_size)
xlabel('E_t[W (\Delta x* ) - W (\Delta x )]', 'FontSize', label_size)

saveas(plot_lines_avgA, graph_path+ "Fig_B3b.png")
saveas(plot_lines_rejA, graph_path+ "Fig_B3a.png")

saveas(plot_lines_avgA, graph_path + "Fig_B3b.eps", 'epsc');
saveas(plot_lines_rejA, graph_path + "Fig_B3a.eps", 'epsc');
